import os
import json
import glob

import time
import math
import cv2
import numpy as np
from tqdm import tqdm


def crop_seg_data(image, mask, crop_h=640, crop_w=640, step_ratio=0.9, padding=False, padding_val=0):
    """
    :param image:待裁剪图
    :param mask: 待裁剪标注二值图
    :param crop_h: 裁剪后的高
    :param crop_w: 裁剪后的宽
    :param step_ratio: 裁剪移动步长比例
    :param padding: True=使用填充；False=不使用填充，取边界的有效区域；当crop尺寸大于图像尺寸时，默认启用padding
    :param padding_val: 填充值（padding=True时生效）
    :return:
        crop_images: 裁剪后的图
        crop_masks: 裁剪后的标注二值图
        crop_positions: [left, top, right, bottom]裁剪图在原图上的位置
    """
    if not (0 < step_ratio <= 1):
        raise ValueError("step_ratio must be between 0 and 1.")
    h, w, c = image.shape
    if crop_h > h or crop_w > w:
        padding = True
    step_x = math.ceil(crop_w * step_ratio)
    step_y = math.ceil(crop_h * step_ratio)

    # 长宽切片数量
    pad_top = pad_bottom = pad_left = pad_right = 0
    if w <= crop_w:
        x_num = 1
        pad_right = crop_w - w
    elif (w - crop_w) % step_x > 0:
        x_num = math.ceil((w - crop_w) / step_x) + 1
        pad_right = step_x - (w - crop_w) % step_x
    else:
        x_num = int((w - crop_w) / step_x)
    if h <= crop_h:
        y_num = 1
        pad_bottom = crop_h - h
    elif (h - crop_h) % step_y > 0:
        y_num = math.ceil((h - crop_h) / step_y) + 1
        pad_bottom = step_y - (h - crop_h) % step_y
    else:
        y_num = int((h - crop_h) / step_y)

    if padding:
        pad_image = cv2.copyMakeBorder(image, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=(padding_val, padding_val, padding_val))
        image = pad_image.copy()
        h, w, c = image.shape

    # Crop
    crop_images = []
    crop_masks = []
    crop_positions = [] # 相对原图的左上/右下角点
    for yn in range(y_num):
        if yn + 1 == y_num and padding is False:
            y_st, y_ed = h - crop_h, h
        else:
            y_st, y_ed = yn * step_y, yn * step_y + crop_h
        for xn in range(x_num):
            if xn + 1 == x_num and padding is False:
                x_st, x_ed = w - crop_w, w
            else:
                x_st, x_ed = xn * step_x, xn * step_x + crop_w

            crop_images.append(image[y_st:y_ed, x_st:x_ed])
            crop_masks.append(mask[y_st:y_ed, x_st:x_ed])
            crop_positions.append((x_st, y_st, x_ed, y_ed))
    return crop_images, crop_masks, crop_positions

def main():
    source_dir = 'D:\\lliujian\\other\\新建文件夹\\测试照片\\测试照片'
    crop_dir = os.path.join(f'{os.path.dirname(source_dir)}', f'{os.path.basename(source_dir)}_crop')

    img_suffix = '.bmp'
    label = 'Needle'
    crop_width, crop_height = 640, 640
    step_ratio = 0.95
    padding = False
    padding_val = 0

    total_time = 0
    if not os.path.exists(crop_dir):
        os.makedirs(crop_dir)
    for i, img_pth in enumerate(tqdm(glob.glob(source_dir + f'\\*{img_suffix}'))):
        img_name = os.path.basename(img_pth)
        _path, _ext = os.path.splitext(img_pth)
        json_path = _path + '.json'

        if os.path.exists(json_path):
            img = cv2.imdecode(np.fromfile(img_pth, dtype=np.uint8), -1)
            h, w, c = img.shape

            with open(json_path, 'r', encoding='utf-8') as f:
                shapes = json.load(f)['shapes']
            polygons = []
            for shape in shapes:
                if shape['shape_type'] == 'polygon' and shape['label'] == label:
                    poly = shape['points']
                    polygons.append(np.array(poly, dtype=np.int32))
                else:
                    print("Skip mismatched label: shape_type={}  label={}".format(shape['shape_type'], shape['label']))
            mask = np.zeros((h, w), np.uint8)
            cv2.fillPoly(mask, polygons, [255])
            # cv2.imencode(img_suffix, mask)[1].tofile(os.path.join(crop_dir, img_name))

            # 裁剪
            start = time.time()
            crop_images, crop_masks, crop_positions =  crop_seg_data(img, mask, crop_height, crop_width, step_ratio, padding, padding_val)
            end = time.time()
            total_time += end - start
            for j, (sub_img, sub_msk, sub_pos) in enumerate(zip(crop_images, crop_masks, crop_positions)):
                sub_img_path = os.path.join(crop_dir, '{}_{}.bmp'.format(img_name.replace(img_suffix, ''), j))
                # 裁剪图
                cv2.imencode(img_suffix, sub_img)[1].tofile(sub_img_path)
                # 裁剪二值图
                # cv2.imencode(img_suffix, sub_msk)[1].tofile(os.path.join(crop_dir, '{}_{}_mask.bmp'.format(img_name.replace(img_suffix, ''), j)))
                # 裁剪标注
                contours, hierarchy = cv2.findContours(sub_msk, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                sub_shapes = []
                for cnt in contours:
                    shp = {
                        "kie_linking": [],
                        "label": label,
                        "score": None,
                        "points": np.squeeze(cnt, axis=1).tolist(),
                        "group_id": None,
                        "description": "",
                        "difficult": False,
                        "shape_type": "polygon",
                        "flags": {},
                        "attributes": {}
                    }
                    sub_shapes.append(shp)
                sub_json_data = {
                    "version": "2.5.2",
                    "flags": {},
                    "shapes": sub_shapes,
                    "imagePath": os.path.basename(sub_img_path),
                    "imageData": None,
                    "imageHeight": sub_img.shape[0],
                    "imageWidth": sub_img.shape[1]
                }
                with open(sub_img_path.replace(img_suffix, '.json'), 'w', encoding='utf-8') as f:
                    json.dump(sub_json_data, f, ensure_ascii=False, indent=4)

            pos_json_data = {
                "name": img_name,
                "image_shape": img.shape,
                "crop_shape": [crop_height, crop_width, c],
                "step_ratio": step_ratio,
                "positions": crop_positions,
                "use_padding": padding,
                "padding_value": padding_val
            }
            with open(os.path.join(crop_dir, img_name.replace(img_suffix, '_positions.json')), 'w', encoding='utf-8') as f:
                json.dump(pos_json_data, f, ensure_ascii=False, indent=4)
    print("mean cost: {:.10f} ms/it".format(total_time / (i + 1) * 1000))

if __name__ == '__main__':
    main()